#/usr/bin/env python
#-*- coding: utf-8 -*-
#
# Authors: Bruno Alves, n.º 6152; Neuza Figueira, n.º 6036
#
# Reference: http://rickyrosario.com/blog/quick-and-dirty-geoip-lookup-function-in-python/

import sys
import re
import GeoIP

from mpl_toolkits.basemap import Basemap
import numpy as np
import matplotlib.pyplot as plt

from collections import OrderedDict

import sqlite3
import csv

import os

from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.lib.units import inch
from reportlab.lib import colors

class Firewall_Data():
	'''Class that allows to process the firewall log (/var/log/ufw.log).''' 

	def read_firewall_log(self):
		'''Method to read the firewall log and obtain data from it.
		   @return info -> Data obtained from the log.
		   @return ip_list -> List of number of attacks by IP.'''
		ip_list = dict()
		info = []
		location = None
		ip = None
		date = None
		time = None

		#file_object = open('/var/log/ufw.log', 'r') 
		file_object = open('ufw.log', 'r') 
		gi = GeoIP.open("GeoIP.dat", GeoIP.GEOIP_STANDARD)
		for line in file_object:
			if not re.search("SRC=192.", line):
				lista = line.split("SRC=")
				if not re.search(":", lista[1].split(' ')[0]):
					ip = lista[1].split(' ')[0]
					date = line[0:6]
					time = line[7:15]

					if ip_list.has_key(ip):
						ip_list[ip] = ip_list[ip] + 1

					else:
						ip_list[ip] = 1
				        
					try:
						location = gi.country_name_by_addr(ip)
					except:
						continue

					finally:
						info.append([ip, date, time, location])

		return info, ip_list


	def geo_ip_lookup(self, ip_address):
		'''Method to seach the coordinates of a given IP address.
		   @param ip_address -> IP to search.
		   @return record_list -> List containing the latitude and longitude.'''
		gi = GeoIP.open("GeoLiteCity.dat", GeoIP.GEOIP_STANDARD)

		gir = gi.record_by_addr(ip_address)

		record_list = []

		if gir is not None:
			record_list.append([gir['city'], gir['latitude'], gir['longitude'], ip_address])

		return record_list


	def build_bars_graph(self, ip_list):
		'''Method to create a simple horizontal bar chart.
		   @param ip_list -> List of all IP's read from the log.
		   @reference: http://matplotlib.org/examples/lines_bars_and_markers/barh_demo.html'''
		new_ip_list = OrderedDict(sorted(ip_list.items(), key=lambda x: x[1]))
		ip = [None] * len(ip_list)
		attempts = [None] * len(ip_list)
		i = 0

		for k in new_ip_list:
			ip[i] = k
			attempts[i] = new_ip_list[k]
			i = i + 1

		y_pos = np.arange(len(ip))

		plt.barh(y_pos, attempts, align='center', color='green')
		plt.yticks(y_pos, ip)
		plt.xlabel('Number of attempts')
		plt.title('Attack Attempts per IP')

		plt.show()


	def build_world_map(self, data):
		'''Method to create a world map chart.
		   @param data -> Data regarding the IP's, suc as latitude and logitude, so their location can be displayed on the map.
		   @reference: http://matplotlib.org/basemap/users/mill.html'''
		m = Basemap(projection='mill',llcrnrlat=-90,urcrnrlat=90,\
		            llcrnrlon=-180,urcrnrlon=180,resolution='c')
		m.drawcoastlines()

		parallels = np.arange(-90.,91.,20.)

		meridians = np.arange(-180.,181.,40.)

		m.fillcontinents(color='#339900',lake_color='#3399FF')
		m.drawmapboundary(fill_color='#3399FF')
		plt.title("Attack Locations")

		for k in range(len(data)):
			try:
				if(str(data[k][0][1]) != str(data[k + 1][0][1])):
					lat, lon = str(data[k][0][1]), str(data[k][0][2]) # Location of Boulder

					xpt,ypt = m(float(lon),float(lat))

					m.plot(xpt,ypt,'ro')
					plt.text(xpt+100000, ypt+100000, str(data[k][0][0]))
			except:
				continue

		plt.show()


	def create_csv_files(self, info):
		'''Method to create the csv file regarding the firewall data.
		   @param info -> Data from the firewall log.'''
		print 'Creating CSV file...'
		outfiledir = os.getcwd() + "/CSV"

		if not os.path.exists(outfiledir):
			os.makedirs(outfiledir)

		name = outfiledir + '/Firewall_Data_Results.csv'

		filename = open(name, 'wb')
		spamWriter = csv.writer(filename, delimiter=',')
		spamWriter.writerow(["IP", "DATE", "TIME", "LOCATION"])

		for i in range(len(info)):
			spamWriter.writerow([info[i][0], info[i][1], info[i][2], info[i][3]])   

		filename.close()
		print 'File successfully created on ' + name 


	def create_database(self, info):
		'''Method to create the sqlite3 database regarding the firewall data.
		   @param info -> Data from the firewall log.'''
		print 'Creating database...'
		outfiledir = os.getcwd() + "/DB"

		if not os.path.exists(outfiledir):
			os.makedirs(outfiledir)

		db_name = outfiledir + "/Firewall_Data_Database.db"
		connection = sqlite3.connect(str(db_name))
		cursor = connection.cursor()

		cursor.execute('''drop table if exists data''')
		connection.commit()

		cursor.execute('''create table if not exists data (id INTEGER PRIMARY KEY,
															ip TEXT,
															date TEXT,
															time TEXT,
															location TEXT)''')

		i = 1

		for j in range(len(info)):
			cmd = "INSERT INTO data VALUES ('{0}','{1}','{2}','{3}','{4}');".format(i,info[j][0], info[j][1], info[j][2], info[j][3])
			cursor.execute(cmd)
			i = i + 1

		#cursor.execute("SELECT * FROM data")
		#rows = cursor.fetchall()
		#for row in rows:
		#	print row

		print 'Database successfully created on ' + db_name


	def create_pdf(self, info):
		'''Method to create the pdf file regarding the firewall data.
		   @param info -> Data from the firewall log.'''
		print 'Creating PDF file...'
		styles = getSampleStyleSheet()

		outfilename = "Firewall_Data.pdf"
		outfiledir = os.getcwd() + "/PDF"
		outfilepath = os.path.join( outfiledir, outfilename )

		if not os.path.exists(outfiledir):
			os.makedirs(outfiledir)

		headerStyle = styles["Heading1"]

		doc = SimpleDocTemplate(outfilepath)
		Story = []
		style = styles["Normal"]
		title = "Firewall Data Log Analysis"
		p = Paragraph(title, styles["Heading1"])

		Story.append(p)
		Story.append(Spacer(1,0.5*inch))

		data = [("IP", "Date", "Time", "Location")]

		for i in range(len(info)):
			data.append([info[i][0], info[i][1], info[i][2], info[i][3]])

		table = Table(data, colWidths=80, rowHeights=30)

		table.setStyle(TableStyle([('INNERGRID', (0,0), (-1, -1), 0.25, colors.black), 
                       ('BOX', (0, 0), (-1, -1), 0.25, colors.black), 
                       ('ALIGN',(0,0),(-1,-1),'CENTER'),
                       ('VALIGN',(0,0),(-1,-1),'MIDDLE')
                     ]))


		Story.append(table)
		doc.build(Story) 
		print 'File successfully created on ' + outfilepath
